#ifndef MULTIGRID_H
#define MULTIGRID_H

#include <eigen3/Eigen/Dense>
#include <eigen3/Eigen/Sparse>
#include <cmath>
#include <vector>
#include <set>
#include <iostream>


enum Interator_Methon { WJ, GS, SOR};
using std::cout;
using std::endl;
using std::vector;
using std::set;
using namespace Eigen;
using Tri = Eigen::Triplet<double>;
using SpMat = Eigen::SparseMatrix<double>;

const double weight = 2.0 /3 ; // Jaccobi迭代和SOR迭代的松弛因子

double L2_Error(VectorXd v1, VectorXd v2) // 返回两个向量之间的L2距离
{
    double error = 0.0;
    for(int i =0; i < v1.size(); i++)
    {
        error += ((v1[i] - v2[i]) * (v1[i] - v2[i]));
    }
    return sqrt(error);
};

/**
 * @brief 对方程Au=F进行一次迭代
 * 
 * @param _A 方程系数矩阵
 * @param _u 迭代初始值 
 * @param _Rhs 方程右端项
 * @param _type 迭代方法
 * @return VectorXd 迭代一次后所得的近似解
 */
VectorXd Interator(SpMat &_A, VectorXd &_u, VectorXd &_Rhs, Interator_Methon _type )
{
    VectorXd v = Eigen::VectorXd::Zero(_Rhs.size());
    if(_type == WJ)
    {
        for( int i =0; i < _u.size(); ++i)
        {
            double temp = _Rhs[i];
            double a_ii = _A.coeffRef(i,i);
            {
                for(Eigen::SparseMatrix<double>::InnerIterator it(_A,i);it;++it)
                {
                    int i_row = it.row();
                    if(i_row == i)
                        continue;
                    else
                        temp -=(it.value() * (_u[i_row]));
                }
                v[i] = temp / a_ii;
            }
        }
        v = (1 - weight) * _u + weight * v ;
    }
    else if (_type == GS)
    {
        v = _u;
        for( int i =0; i < v.size(); ++i)
        {
            double temp = _Rhs[i];
            double a_ii = _A.coeffRef(i,i);
            {
                for(Eigen::SparseMatrix<double>::InnerIterator it(_A,i);it;++it)
                {
                    int i_row = it.row();
                    if(i_row == i)
                        continue;
                    else
                        temp -=(it.value() * (v[i_row]));
                }
                v[i] = temp / a_ii;
            }
        }
    }
    else if(_type == SOR)
    {
        v = _u;
        for( int i =0; i < v.size(); ++i)
        {
            double temp = _Rhs[i];
            double a_ii = _A.coeffRef(i,i);
            {
                for(Eigen::SparseMatrix<double>::InnerIterator it(_A,i);it;++it)
                {
                    int i_row = it.row();
                    if(i_row == i)
                        continue;
                    else
                        temp -=(it.value() * (v[i_row]));
                }
                v[i] = temp / a_ii * weight + v[i] * (1 - weight);
            }
        }
    }
    return v ;
};

void DealWithBoundary(SpMat &A, VectorXd &Rhs, set<size_t> &boundary_index)
{
    int n_col = Rhs.size();
    VectorXd u = VectorXd::Zero(n_col);
    for(int i =0; i < n_col; i++)
    {
        if(boundary_index.find(i) != boundary_index.end())      //第i个是边界，则Rhs[i]/A(i,i)就是真解
        {
            u[i] = Rhs[i] / A.coeffRef(i,i);
        }
    }
    for(int i =0; i <n_col; i++)
    {
        if(boundary_index.find(i) != boundary_index.end())
        {
            for(Eigen::SparseMatrix<double>::InnerIterator it(A,i);it;++it)
            {
                int row = it.row();     //注意，以下所有的取A（i，j）的值需要保证该值存在，不然会使迭代器失效
                if(row == i)
                    continue;
                Rhs[row] -= A.coeffRef(row, i) * u[i];
                A.coeffRef(i,row) = 0.0;
                A.coeffRef(row,i) = 0.0;
            }
        }
    }
    for(int i =0; i < n_col; i++)
    {
        if(boundary_index.find(i) != boundary_index.end())      
        {
            Rhs[i] = u[i] *A.coeffRef(i,i);
        }
    }
};


class Multigrid_Twogrid
{
protected:
    SpMat A_fine ;          //细网格下的矩阵
    SpMat A_coarser;         //粗网格下的矩阵
    VectorXd Rhs_fine;      //细网格下的方程右端项
    VectorXd Rhs_coarser;    //粗网格下的方程右端项（残差项）
    set<size_t> fine_boundary;   //细网格下的边界点
    set<size_t> coarser_boundary;    //粗网格下的边界点
    double Tolerance = 1e-8;   //粗网格求解使用的迭代法的终止精度
    vector<VectorXd> solution_u;    //u每次迭代后的近似解
    size_t iter_num;        //迭代终止时整个过程所迭代的次数
    bool is_convergence = false;    //判断是否已解出方程
    Interator_Methon interator_type = GS ;  //迭代的方法
public:
    void set_fine_boundary(set<size_t> &_fine_boundary) { fine_boundary = _fine_boundary; };
    void set_coarser_boundary(set<size_t> &_coarser_boundary) { coarser_boundary = _coarser_boundary; };
    void set_fineMat(SpMat &_fine_mat) { A_fine = _fine_mat; };
    void set_coarserMat(SpMat &_coarser_mat) { A_coarser = _coarser_mat; };
    void set_rhs_fine(VectorXd &_rhs_fine) { Rhs_fine = _rhs_fine; };
    void set_rhs_coarser(VectorXd &_rhs_coarser) { Rhs_coarser = _rhs_coarser; };
    void set_tolerance(double _tolerance) { Tolerance = _tolerance; };
    Multigrid_Twogrid() = default;
    Multigrid_Twogrid(SpMat &_A_fine, SpMat &_A_coarser, VectorXd &_Rhs_fine, VectorXd &_Rhs_coarser, set<size_t> &_fine_boundary, set<size_t> &_coarser_boundary):
            A_fine(_A_fine), A_coarser(_A_coarser), Rhs_fine(_Rhs_fine), Rhs_coarser(_Rhs_coarser), fine_boundary(_fine_boundary), coarser_boundary(_coarser_boundary), iter_num(0) { };
    virtual VectorXd interpolation(VectorXd &u_2h) = 0 ;    //将粗网格下的值差值到细网格中
    virtual VectorXd restriction(VectorXd &u_h) = 0 ;       //将细网格下的值投影到粗网格中
    void compute();         //开始计算
    void set_interator_type(Interator_Methon _interator_type ) { interator_type = _interator_type; };
};

void Multigrid_Twogrid::compute()
{   
    int vis1 = 5, vis2 = 5;
    VectorXd u = VectorXd::Zero(Rhs_fine.size());
    VectorXd u_iter1 = Interator(A_fine, u, Rhs_fine, interator_type) ;
    while(!is_convergence)
    {
        for(int i =0; i < vis1; i++) //先在细网格上迭代vis1次
        {
            iter_num ++ ;
            solution_u.push_back(u_iter1);
            u_iter1 = Interator(A_fine, u, Rhs_fine, interator_type) ;
            if(L2_Error(u_iter1,u) < Tolerance)
            {
                cout <<"calculate over!\n" << "The number of interator is :" << iter_num << endl;
                is_convergence = true;
                break;
            }
            u = u_iter1;
        }
        if(is_convergence)
            break ;
        VectorXd r_h = Rhs_fine - A_fine*u;
        VectorXd r_2h = restriction(r_h);
        set_rhs_coarser(r_2h);
        VectorXd e_2h = VectorXd::Zero(r_2h.size());
        VectorXd e_2h_iter1 = Interator(A_coarser, e_2h, r_2h, interator_type);
        for(int i =0; i < vis2; i++)
        {
            if(L2_Error(e_2h_iter1, e_2h) < Tolerance)
            {
                break;
            };
            e_2h = e_2h_iter1;
            e_2h_iter1 = Interator(A_coarser, e_2h, r_2h, interator_type);
            iter_num ++;    
        }
        VectorXd e_h = interpolation(e_2h);
        u = u + e_h ;
        u_iter1 = Interator(A_fine, u, Rhs_fine, interator_type) ;
        iter_num ++;
        if(L2_Error(u_iter1,u) < Tolerance)
        {
            cout <<"calculate over!\n" << "The number of interator is :" << iter_num << endl;
            is_convergence = true;
            break;
        }
    }
};


class DIM1_Multigrid_Twogrid : public Multigrid_Twogrid
{
public:
    DIM1_Multigrid_Twogrid() = delete;
    DIM1_Multigrid_Twogrid(SpMat &_A_fine, SpMat &_A_coarser, VectorXd &_Rhs_fine, VectorXd &_Rhs_coarser, set<size_t> &_fine_boundary, set<size_t> &_coarser_boundary):
            Multigrid_Twogrid(_A_fine, _A_coarser, _Rhs_fine, _Rhs_coarser, _fine_boundary, _coarser_boundary){ };
    VectorXd interpolation(VectorXd &u_2h);
    VectorXd restriction(VectorXd &u_h);
};

VectorXd DIM1_Multigrid_Twogrid::interpolation(VectorXd &u_2h)
{
    int n_size = u_2h.size();
    VectorXd u_h = VectorXd::Zero(2*n_size -1);
    for(int j =0; j < 2*n_size -1; j++)
    {
        if(j %2 == 0)
            u_h[j] = u_2h[j/2];
        else
            u_h[j] = (u_2h[j/2] + u_2h[j/2 + 1]) / 2;
    }
    return u_h;
};

VectorXd DIM1_Multigrid_Twogrid::restriction(VectorXd &u_h)
{
    int u_h_size = u_h.size();
    int u_2h_size = u_h_size / 2 + 1 ;
    VectorXd u_2h = VectorXd::Zero(u_2h_size);
    u_2h[0] = u_h[0];
    u_2h[u_2h_size - 1] = u_h[u_h_size - 1];
    for(int i = 1; i < u_2h_size - 1; i++)
        u_2h[i] = (u_h[2*i-1] + 2* u_h[2 *i] + u_h[2*i+1]) / 4;
    return u_2h;
};
/**
 * @brief 在这里所使用的网格以及为有限元方法生成的Q1网格
 * 
 */
class DIM2_Multigrid_Twogrid : public Multigrid_Twogrid
{
private:
    int nx_2h;  //粗网格下x方向上的单元数
    int ny_2h;  //粗网格下y方向上的单元数
public:
    DIM2_Multigrid_Twogrid() = delete;
    DIM2_Multigrid_Twogrid(SpMat &_A_fine, SpMat &_A_coarser, VectorXd &_Rhs_fine, VectorXd &_Rhs_coarser, set<size_t> &_fine_boundary, 
                        set<size_t> &_coarser_boundary, int _nx_2h, int _ny_2h):
                    Multigrid_Twogrid(_A_fine, _A_coarser, _Rhs_fine, _Rhs_coarser, _fine_boundary, _coarser_boundary), nx_2h(_nx_2h), ny_2h(_ny_2h)     { };
    VectorXd interpolation(VectorXd &u_2h);
    VectorXd restriction(VectorXd &u_h);
    size_t coarser_index(size_t _i, size_t _j);
    size_t fine_index(size_t _i, size_t _j);
};

inline size_t DIM2_Multigrid_Twogrid::coarser_index(size_t _i, size_t _j)
{
    return (_j + _i *(nx_2h + 1));
};

inline size_t DIM2_Multigrid_Twogrid::fine_index(size_t _i, size_t _j)
{
    return (_j + _i *(2 *nx_2h + 1));
};
VectorXd DIM2_Multigrid_Twogrid::interpolation(VectorXd &u_2h)
{
    VectorXd u_h = VectorXd::Zero(Rhs_fine.size());
    for(size_t i =0; i <= nx_2h; i++)
        for(size_t j =0; j <= ny_2h; j++)
        {
            u_h[fine_index(2*i,2*j)] = u_2h[coarser_index(i,j)];
        }
    for(size_t i =0; i < nx_2h; i++)
        for(size_t j =0; j < ny_2h; j++)
        {
            u_h[fine_index(2*i+1,2*j)] = (u_2h[coarser_index(i,j)] + u_2h[coarser_index(i+1,j)]) / 2;
            u_h[fine_index(2*i,2*j+1)] = (u_2h[coarser_index(i,j)] + u_2h[coarser_index(i,j+1)]) / 2;
            u_h[fine_index(2*i+1,2*j+1)] = (u_2h[coarser_index(i,j)] + u_2h[coarser_index(i+1,j)] + u_2h[coarser_index(i,j+1)] + u_2h[coarser_index(i+1,j+1)]) / 4;
        }
    return u_h;
};

VectorXd DIM2_Multigrid_Twogrid::restriction(VectorXd &u_h)
{
    VectorXd u_2h = VectorXd::Zero((nx_2h+1)*(ny_2h+1));
    for(size_t i =0; i <= nx_2h; i++)
        for(size_t j =0; j <= ny_2h; j++)
        {
            if(coarser_boundary.find(coarser_index(i,j)) != coarser_boundary.end())     //当该点为边界点时
                u_2h[coarser_index(i,j)] = u_h[fine_index(2*i,2*j)];
            else{
                u_2h[coarser_index(i,j)] = 1.0 / 16 *(u_h[fine_index(2*i-1,2*j-1)] + u_h[fine_index(2*i-1,2*j+1)] + u_h[fine_index(2*i+1,2*j-1)] + u_h[fine_index(2*i+1,2*j+1)]
                        + 2 *(u_h[fine_index(2*i,2*j-1)] + u_h[fine_index(2*i,2*j+1)] + u_h[fine_index(2*i-1,2*j)] + u_h[fine_index(2*i+1,2*j)]) + 4 *u_h[fine_index(2*i,2*j)]);
            }
        }
    return u_2h;
};

#endif